{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5dd31051",
   "metadata": {},
   "source": [
    "LLM 模型的权重需要大量的内存，但是通过量化，我们可以将非常大的模型压缩到可以在普通硬件上运行的大小。\n",
    "\n",
    "在这节课里，可以学习量化的工作原理，并实现一个算法\n",
    "\n",
    "将一个模型进行 int8 量化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b3b9965",
   "metadata": {},
   "outputs": [],
   "source": [
    "# step1 导入相关包\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import time\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from transformers.models.gpt2.modeling_gpt2 import GPT2Model\n",
    "# 从之前的课程中导入我们的生成函数\n",
    "from utils import generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86049a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# step2 加载 model 和 tokenizer\n",
    "model_name ='gpt2'\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7cc49b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 连续批处理技术，也是在 批处理，所以我们需要引入填充标记\n",
    "# Define PAD token as EOS token\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "model.config.pad_token_id = model.config.eos_token_id\n",
    "\n",
    "# 定义填充标记放置的位置  left or  right\n",
    "tokenizer.padding_side = \"left\"  # 或 \"right\"\n",
    "tokenizer.truncation_side = \"left\"  # 或 \"right\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56ceab38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 零点量化\n",
    "# 编写一个函数来帮助我们在代码中执行这个量化步骤\n",
    "# 我们要做的第一件事是计算张量的最小值和最大值,然后我们要计算两组元数据值,一个是比例尺(比例因子),一个是零点。\n",
    "\n",
    "def quantize(t):\n",
    "    # obtain range of values in the tensor to map between @ and 255\n",
    "    min_val, max_val = t.min(), t.max()\n",
    "\n",
    "    # determine the \"zero-point\", or value in the tensor to map to @\n",
    "    scale = (max_val - min_val) / 255\n",
    "    zero_point = min_val\n",
    "\n",
    "    # quantize and clamp to ensure we're in [0, 255]\n",
    "    t_quant = (t - zero_point) / scale\n",
    "    t_quant = torch.clamp(t_quant, min=0, max=255)\n",
    "\n",
    "    # keep track of scale and zero_point for reversing quantization\n",
    "    state = (scale, zero_point)\n",
    "\n",
    "    # cast to uint8 and return\n",
    "    t_quant = t_quant.type(torch.uint8)\n",
    "    return t_quant, state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8129a942",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 让我们运行一个快速测试,看看这个量化函数对我们模型中的随机张量有什么影响。\n",
    "# 所以我们将从第一个注意力块层中获取一个张量,并打印出 它的形状和它的一小部分值。\n",
    "t = model. transformer.h[0].attn.c_attn.weight.data\n",
    "print(t, t.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb41933",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 现在让我们使用这个张量t 调用我们的量化函数,并打印出量化状态中的值的样本,以及tq中的最小值和最大值,如果一切按我们预期的那样工作,它们应该对应于0和255。\n",
    "\n",
    "t_q, state = quantize(t)\n",
    "print(t_q, t_q.min(), t_q.max())\n",
    "\n",
    "# 打印看输出结果，可以看到 量化函数达到了我们的预期效果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc66b018",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 接下来实现 一个函数来反转量化步骤, 这应该会简单得多。\n",
    "def dequantize(t, state):\n",
    "    scale, zero_point = state\n",
    "    return t.to(torch.float32) * scale + zero_point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a71716b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "t_rev = dequantize(t_q, state)\n",
    "print(t_rev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b66b2a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.abs(t - t_rev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "169575c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 接下来，我们把这种量化技术应用到整个模型上\n",
    "# 先看下：模型对于特定请求的输出是什么？ 好去跟量化的结果做比较\n",
    "response_expected = generate(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    [(\"The quick brown fox jumped over the\", 10)]\n",
    ") [0]\n",
    "response_expected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89261374",
   "metadata": {},
   "outputs": [],
   "source": [
    "def quantize_model(model):\n",
    "    states = {}\n",
    "\n",
    "    for name, param in model.named_parameters():\n",
    "        param.requires_grad = False\n",
    "        param.data, state = quantize(param.data)\n",
    "        states[name] = state\n",
    "    return model, states\n",
    "\n",
    "quant_model, states = quantize_model(model)\n",
    "\n",
    "quant_model.get_memory_footprint() # 137022768"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a3ed5fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def size_in_bytes(t):\n",
    "    return t.numel() * t.element_size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ec91380",
   "metadata": {},
   "outputs": [],
   "source": [
    "sum([\n",
    "    size_in_bytes(v[0]) + size_in_bytes(v[1])\n",
    "    for v in sta[es.values()\n",
    "])\n",
    "    \n",
    "# 1184"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8def9d32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 编写dequantize_model\n",
    "def dequantize_model(model, states):\n",
    "    for name, param in model.named_parameters():\n",
    "        state = states[name]\n",
    "        param.data = dequantize(param.data, state)\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70e8bd76",
   "metadata": {},
   "outputs": [],
   "source": [
    "dequant_model = dequantize_model(quant_model, states)\n",
    "\n",
    "dequant_model.get_memory_footprint() # 510342192"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9197df2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 我们来使用 dequant_model 模型，而不是原始模型: model\n",
    "response_expected = generate(\n",
    "    dequant_model,\n",
    "    tokenizer,\n",
    "    [(\"The quick brown fox jumped over the\", 10)]\n",
    ") [0]\n",
    "response_expected\n",
    "\n",
    "# 打印输出结果，输出仍然是合理的，但是会有点重复，语义上仍然是正确的"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81723e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fix_dtype_post_quantization_to \"pretend\" to be fp32\n",
    "def get_float32_dtype(self):\n",
    "    return torch.float32\n",
    "\n",
    "GPT2Model.dtype = property(get_float32_dtype)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
